#!/usr/bin/env python3
# SPDX-License-Identifier: LGPL-2.1-or-later

import logging
import sys
import tempfile
from pathlib import Path
from typing import Optional

from mkosi import identify_cpu
from mkosi.archive import make_cpio
from mkosi.config import OutputFormat
from mkosi.initrd import KernelInstallContext
from mkosi.log import log_setup
from mkosi.run import run, uncaught_exception_handler
from mkosi.sandbox import umask
from mkosi.util import PathString


def we_are_wanted(context: KernelInstallContext) -> bool:
    return context.uki_generator == "mkosi" or context.initrd_generator in ("mkosi", "mkosi-initrd")


def build_microcode_initrd(output: Path) -> Optional[Path]:
    vendor, ucode = identify_cpu(Path("/"))

    if vendor is None:
        logging.warning("Unable to determine the vendor of your CPU, not adding microcode")
        return None

    if ucode is None:
        logging.warning("Unable to find microcode for your CPU in /usr/lib/firmware, not adding microcode")
        return None

    with tempfile.TemporaryDirectory() as tmp:
        root = Path(tmp) / "initrd-microcode-root"
        destdir = root / "kernel/x86/microcode"

        with umask(~0o755):
            destdir.mkdir(parents=True, exist_ok=True)

            with (destdir / f"{vendor}.bin").open("wb") as f:
                f.write(ucode.read_bytes())

        make_cpio(root, output)

    return output


@uncaught_exception_handler()
def main() -> None:
    context = KernelInstallContext.parse(
        name="50-mkosi.install",
        description="kernel-install plugin to build initrds or Unified Kernel Images using mkosi",
    )

    log_setup(default_log_level="info" if context.verbose else "warning")

    if context.command != "add" or not we_are_wanted(context):
        logging.info("mkosi-initrd is not enabled, skipping")
        return

    # If kernel-install was passed a UKI, there's no need to build anything ourselves.
    if context.image_type == "uki":
        logging.info("Provided kernel image is already a unified kernel image, skipping mkosi-initrd")
        return

    # If the initrd was provided on the kernel command line, we shouldn't generate our own.
    if context.layout != "uki" and context.initrds:
        logging.info("Pre-built initrds were provided, skipping mkosi-initrd")
        return

    if context.layout == "uki" and context.uki_generator == "mkosi":
        format = OutputFormat.uki
    else:
        format = OutputFormat.cpio

    output = "initrd" if format == OutputFormat.cpio else "uki.efi"

    cmdline: list[PathString] = [
        "mkosi-initrd",
        "--kernel-version", context.kernel_version,
        "--format", str(format),
        "--output", output,
        "--output-dir", context.staging_area,
        "--kernel-image", context.kernel_image,
    ]  # fmt: skip

    if context.verbose:
        cmdline += ["--debug"]

    logging.info(f"Building {output}")

    run(cmdline, stdin=sys.stdin, stdout=sys.stdout)

    if format == OutputFormat.cpio:
        build_microcode_initrd(context.staging_area / "microcode")


if __name__ == "__main__":
    main()
